from collections.abc import Iterable, Mapping, Sequence
from typing import TypeVar

import numpy as np

K = TypeVar("K")
T = TypeVar("T")


def random_split(
    items: Iterable[T],
    weights: Mapping[K, float] | Iterable[float], *,
    random_seed: int = 42,
) -> Mapping[K, tuple[T, ...]] | tuple[tuple[T, ...], ...]:
    items = items if isinstance(items, Sequence) else tuple(items)
    if isinstance(weights, Mapping):
        keys, weights = zip(*weights.items())
        indices = shuffle_and_split_indices(len(items), weights, random_seed=random_seed)
        return {k: tuple(items[i] for i in indices) for k, indices in zip(keys, indices)}
    if isinstance(weights, Iterable):
        weights = tuple(weights)
        indices = shuffle_and_split_indices(len(items), weights, random_seed=random_seed)
        return tuple(tuple(items[i] for i in indices) for indices in indices)
    raise TypeError(f"Expected weights to be Mapping or Iterable, got {type(weights)}")


def shuffle_and_split_indices(n: int, weights: Iterable[float], *, random_seed: int = 42):
    weights = np.asarray(tuple(weights), dtype=np.float32)
    counts = n * weights / weights.sum()
    tails = np.cumsum(counts)
    tails = np.asarray(np.round(tails), dtype=np.int64)
    rng = np.random.default_rng(random_seed)
    indices = rng.permutation(n)
    return np.split(indices, tails[:-1])
